from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim
from configs import BGE_BASE_MODEL_PATH, device
from tqdm import trange
import numpy as np
import torch
import logging
from torch import Tensor
# BGE_BASE_MODEL_PATH = "../embedding_models/bge-base"
logger = logging.getLogger(__name__)
# device = torch.device("cuda:3")

show_progress_bar = (logger.getEffectiveLevel()==logging.INFO or logger.getEffectiveLevel()==logging.DEBUG)
model = SentenceTransformer(BGE_BASE_MODEL_PATH)
model.eval()
model = model.to(device)


def batch_to_device(batch):
    for key in batch:
        if isinstance(batch[key], Tensor):
            batch[key] = batch[key].to(device)
    return batch


def bge_encode(text_list, is_normalize=True):
    if not isinstance(text_list, list):
        text_list = [text_list]
        flag = 0
    else:
        flag = 1
    batch_size = 32
    length_sorted_idx = np.argsort([-len(sen) for sen in text_list])
    sentences_sorted = [text_list[idx] for idx in length_sorted_idx]
    all_embeddings = []
    with torch.no_grad():
        for start_index in trange(0, len(text_list), batch_size, desc="Batches", disable=not show_progress_bar):
            sentences_batch = sentences_sorted[start_index:start_index+batch_size]
            encoded_input = model.tokenize(sentences_batch)
            encoded_input = batch_to_device(encoded_input)
            fea = model(encoded_input)
            embeddings = fea["sentence_embedding"]
            embeddings = embeddings.detach()
            if is_normalize:
                embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
            embeddings = embeddings.cpu()
            all_embeddings.extend(embeddings)
    all_embeddings = [all_embeddings[idx] for idx in np.argsort(length_sorted_idx)]
    all_embeddings = np.asarray([emb.numpy() for emb in all_embeddings])
    if flag == 0:
        return all_embeddings[0]
    else:
        return all_embeddings
    return all_embeddings


# embed = bge_encode(["1.1合同构成本合同是您与我们约定保险权利义务关系的协议，包括本保险条款、保险单或者其他保险凭证、投保单、与本合同有关的投保文件、合法有效(见 8.1)的声明、批注、批单及其他您与我们共同认可的书面协议。1.2合同成立及生效您提出保险申请且我们同意承保，本合同成立。本合同的成立日、生效日以保险单记载的日期为准。1.3投保年龄投保年龄指您投保时被保险人的年龄，以周岁(见 8.2)计算。首次投保（见 8.3）或重新投保（见 8.4）的投保年龄范围为出生 28 天（含）至 60 周岁（含）。1.4投保人范围除另有约定外，投保时已参加宁波市城镇职工基本医疗保险或宁波市城乡居民基本医疗保险的人员可作为本保险的投保人。1.5被保险人范围除另有约定外，投保时已参加宁波市城镇职工基本医疗保险或宁波市城乡居民基本医疗保险、身体健康且符合本公司承保条件者均可作为被保险人参加本保险。被保险人应为投保人本人或其配偶、子女、父母等具有保险利益的人。", "2.1保险金额本合同的保险金额由您在投保时与我们约定，并在保险单上载明。2.2保险期间本合同的保险期间为 1 年，自本合同生效日零时开始，至保险单载明的保险期间期满日的 24 时止。若您与我们按照本条款 4.3 约定进行续保，续保合同自本合同期满日次日零时起生效，有效期为 1 年，至续保合同的保险期间期满日 24 时止。续保合同，指按本条款约定办理续保而生效的保险合同。2.3等待期您为被保险人首次投保本合同或者重新投保本合同时，被保险人自本合同生效之日起 90 日内经医院(见 8.5)初次确诊(见 8.6)非因意外伤害（见 8.7）罹患本合同定义的重大疾病，我们不承担给付重大疾病保险金的责任，但向您退还已交纳的本合同的保险费，本合同终止。这 90 日称为保险责任等待期。续保合同无等待期。本合同定义的重大疾病载明于本合同“7 重大疾病定义”。2.4保险责任在本合同保险期间内，我们承担下列保险责任：重大疾病保险金被保险人于本合同生效后经医院初次确诊因意外伤害导致罹患本合同定义的重大疾病（无论一种或多种），或者于等待期后经医院初次确诊非因意外伤害导致罹患本合同定义的重大疾病（无论一种或者多种），我们按照本合同约定的保险金额向重大疾病保险金受益人给付重大疾病保险金，本合同终止。2.5责任免除因下列情形之一导致被保险人发生本合同定义的重大疾病的，我们不承担给付重大疾病保险金的责任：（1）投保人对被保险人的故意杀害、故意伤害；（2）被保险人故意犯罪或者抗拒依法采取的刑事强制措施；（3）被保险人自杀或故意自伤，但被保险人自杀或故意自伤时为无民事行为能力人的除外；（4）被保险人主动吸食或者注射毒品（见 8.8）；（5）被保险人酒后驾驶（见 8.9）、无合法有效驾驶证驾驶（见 8.10），或者驾驶无合法有效行驶证（见 8.11）的机动车（见 8.12）；（6）被保险人感染艾滋病病毒或者患艾滋病(见 8.13)；（7）战争、军事冲突、暴乱或者武装叛乱、恐怖活动；（8）核爆炸、核辐射或者核污染；（9）遗传性疾病(见 8.14)，先天性畸形、变形或者染色体异常(见 8.15)。因上述第（1）项情形导致被保险人发生本合同定义的重大疾病的，本合同终止，我们向被保险人给付本合同终止时的现金价值(见 8.16)；因上述第（2）至第（9）项中任一情形导致被保险人发生本合同定义的重大疾病的，本合同终止，我们向您退还本合同终止时的现金价值。", "3.1受益人除另有指定外，本合同的受益人为被保险人本人。3.2保险事故通知您、被保险人或者受益人知道保险事故发生后应当在 10 日内通知我们。故意或者因重大过失未及时通知，致使保险事故的性质、原因、损失程度等难以确定的，我们对无法确定的部分不承担给付保险金的责任，但我们通过其他途径已经及时知道或者应当及时知道保险事故发生或者虽未及时通知但不影响我们确定保险事故的性质、原因、损失程度的除外。若您、被保险人或者受益人知道保险事故发生后，未及时通知我们，且您要求解除本合同导致我们按本合同约定向您退还现金价值或保险费后，我们按本合同约定需给付保险金时，有权扣减已给付的现金价值或保险费。3.3保险金申请在申请重大疾病保险金时，受益人作为申请人须填写领取保险金申请书，并须提供下列证明和资料的原件：（1）本合同；（2）受益人的有效身份证件（见 8.17）；（3）由医院专科医生（见 8.18)出具的被保险人的疾病诊断证明书，以及由医院出具的与该疾病诊断证明书相关的病理显微镜检查、血液检验及其他科学方法检验报告；（4）所能提供的与确认保险事故的性质、原因等有关的其他证明和资料。特别注意事项以上证明和资料不完整的，我们将及时一次性通知受益人补充提供有关的证明和资料。委托他人领取保险金时，受托人还必须提供本人的有效身份证件及委托人亲笔签名的授权委托书。保险金作为被保险人遗产时，继承人还必须提供可证明其具备合法继承权的相关权利文件。受益人或者继承人为无民事行为能力人或者限制民事行为能力人时，由其合法监护人代其申请领取保险金，其合法监护人还必须提供受益人或者继承人为无民事行为能力人或者限制民事行为能力人的证明和监护人具有合法监护权的证明。3.4保险金给付我们在收到领取保险金申请书及本合同约定的证明和资料后，将在 5 日内作出核定；情形复杂的，在 30 日内作出核定。对属于保险责任的，我们在与受益人达成给付保险金的协议后 10 日内，履行给付保险金义务。我们未及时履行前款规定的义务的，对属于保险责任的，除支付保险金外，应当赔偿受益人因此受到的利息损失。利息按照我们确定的利率按单利计算，且我们确定的利率不低于中国人民银行公布的金融机构人民币活期存款基准利率。对不属于保险责任的，我们自作出核定之日起 3 日内向受益人发出拒绝给付保险金通知书并说明理由。我们在收到领取保险金申请书及有关证明和资料之日起 60 日内，对给付保险金的数额不能确定的，根据已有证明和资料可以确定的数额先予支付；我们最终确定给付保险金的数额后，将支付相应的差额。"])
# print(embed.shape)


